#!/usr/bin/python -S
"""
os_process.py

A small wrapper for Python's subprocess.  Calling out to subprocesses is I/O, so
it's useful to have this interface for mocking in tests.  

This simplifies basic invocations of subprocess, and also has some
cross-platform fixes.
"""

__author__ = 'Andy Chu'


import sys
import subprocess

import log
import records


class Error(Exception):
  pass


class Runner(object):
  """An interface for retrieving various data from a subprocess run.

  TODO:
    - Cross platform way to invoke the Python interpreter?
  """

  def __init__(self, log=log.NO_LOG, check_exit_code=False, **kwargs):
    """
    Args:
      log: Optional logger
      check_exit_code: If true, insist on a non-zero exit code
      kwargs: Passed through to subprocess.Popen
    """
    self.log = log
    self.check_exit_code = check_exit_code
    # Change the default
    if 'universal_newlines' not in kwargs:
      kwargs['universal_newlines'] = True
    self.kwargs = kwargs

  def _MakeArgs(self, cmd):
    kwargs = self.kwargs

    if isinstance(cmd, list):
      kwargs['shell'] = False
      # Run with python.exe on Windows, to work around stupid bugs
      if sys.platform == 'win32' and cmd and cmd[0].endswith('.py'):
        cmd = [sys.executable] + cmd  # Don't mutate it
    elif isinstance(cmd, basestring):
      kwargs['shell'] = True
    else:
      raise AssertionError('Command %r must be list or string' % cmd)

    return cmd, kwargs

  def _LogRun(self, cmd):
    self.log.info('\tRunning process %s', cmd)

  def _Check(self, exit_code, cmd):
    if self.check_exit_code and exit_code != 0:
      raise Error('Command %r exited with code %s' % (cmd, exit_code))

  def _ErrorWrapper(self, func, cmd, **kwargs):
    """OSError isn't very informative, so raise a clearer error."""
    try:
      return func(cmd, **kwargs)
    except OSError, e:
      raise Error('Error executing %s: %s' % (cmd, e))

  def StdoutLines(self, cmd):
    """
    Return an iterator over stdout lines which automatically closes after the
    first iteration.
    """
    self._LogRun(cmd)
    cmd, kwargs = self._MakeArgs(cmd)
    p = self._ErrorWrapper(
        subprocess.Popen, cmd, stdout=subprocess.PIPE, **kwargs)
    def Iter():
      for line in p.stdout:
        yield line
      p.stdout.close()
      exit_code = p.wait()
      self._Check(exit_code, cmd)
    return Iter()

  def Run(self, cmd):
    """Just run the command, without capturing anything.

    Raises:
      Error: If the command returns a non-zero exit code.
    """
    self._LogRun(cmd)
    cmd, kwargs = self._MakeArgs(cmd)
    exit_code = self._ErrorWrapper(subprocess.call, cmd, **kwargs)
    self._Check(exit_code, cmd)

  def Result(self, cmd):
    """Quick and dirty way to get the result when we don't care about streaming
    the output.

    Returns:
      A record with stdout, stderr, and exit_code fields.

    Raises:
      Error: If the command returns a non-zero exit code.
    """
    self._LogRun(cmd)
    cmd, kwargs = self._MakeArgs(cmd)
    p = self.Pipes(cmd, 'OE')
    stdout = p.stdout.read()
    stderr = p.stderr.read()
    exit_code = p.wait()
    self._Check(exit_code, cmd)
    return records.Record(stdout=stdout, stderr=stderr, exit_code=exit_code)

  def Pipes(self, cmd, descriptors='OE'):
    """Open a pipes to various streams.  Gives greater control.

    Args:
      cmd: Either an argument list or a shell string.
      descriptors: A string containing some of 'OEI', for stdout, stderr, or
          stdin.  subprocess will open pipes to these descriptors.

    Returns:
      A subprocess.Popen instance
    """
    self._LogRun(cmd)
    cmd, kwargs = self._MakeArgs(cmd)

    for char in descriptors:
      if char == 'O':
        kwargs['stdout'] = subprocess.PIPE
      elif char == 'E':
        kwargs['stderr'] = subprocess.PIPE
      elif char == 'I':
        kwargs['stdin'] = subprocess.PIPE
      else:
        raise Error(
            'Invalid descriptor character %r in %r' % (char, descriptors))
    return self._ErrorWrapper(subprocess.Popen, cmd, **kwargs)
